import numpy as np
import struct


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def ivecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('i' * len(x), *x))
    f.close()


if __name__ == '__main__':
    dataset_name_l = ['text-to-image', 'tiny5m']
    basic_dir = '/home/zhengbian/mips-graph/data'
    for ds_name in dataset_name_l:
        acc_gnd = None
        for i in range(10):
            tmp_dir = '%s/%s_bias/part_data/cosine_gnd_idx_%d.ivecs' % (basic_dir, ds_name, i + 1)
            tmp_gnd, _ = ivecs_read(tmp_dir)
            if acc_gnd is not None:
                acc_gnd = np.append(acc_gnd, tmp_gnd, axis=0)
            else:
                acc_gnd = tmp_gnd
        res_fname = '%s/%s_bias/cosine_gnd_idx.ivecs' % (basic_dir, ds_name)
        ivecs_write(res_fname, acc_gnd)
        print("write cosine_gnd dataset %s", ds_name)

        acc_gnd = None
        for i in range(10):
            tmp_dir = '%s/%s_bias/part_data/ip_gnd_idx_%d.ivecs' % (basic_dir, ds_name, i + 1)
            tmp_gnd, _ = ivecs_read(tmp_dir)
            if acc_gnd is not None:
                acc_gnd = np.append(acc_gnd, tmp_gnd, axis=0)
            else:
                acc_gnd = tmp_gnd
        res_fname = '%s/%s_bias/ip_gnd_idx.ivecs' % (basic_dir, ds_name)
        ivecs_write(res_fname, acc_gnd)
        print("write ip_gnd dataset %s", ds_name)

        acc_gnd = None
        for i in range(10):
            tmp_dir = '%s/%s_bias/part_data/l2_gnd_idx_%d.ivecs' % (basic_dir, ds_name, i + 1)
            tmp_gnd, _ = ivecs_read(tmp_dir)
            if acc_gnd is not None:
                acc_gnd = np.append(acc_gnd, tmp_gnd, axis=0)
            else:
                acc_gnd = tmp_gnd
        res_fname = '%s/%s_bias/l2_gnd_idx.ivecs' % (basic_dir, ds_name)
        ivecs_write(res_fname, acc_gnd)
        print("write l2_gnd dataset %s", ds_name)

        acc_gnd = None
        for i in range(10):
            tmp_dir = '%s/%s_bias/part_data/norm_gnd_idx_%d.ivecs' % (basic_dir, ds_name, i + 1)
            tmp_gnd, _ = ivecs_read(tmp_dir)
            if acc_gnd is not None:
                acc_gnd = np.append(acc_gnd, tmp_gnd, axis=0)
            else:
                acc_gnd = tmp_gnd
        res_fname = '%s/%s_bias/norm_gnd_idx.ivecs' % (basic_dir, ds_name)
        ivecs_write(res_fname, acc_gnd)
        print("write norm_gnd dataset %s", ds_name)
